"""
ipfs 操作相关
"""

import io
import torch
import ipfshttpclient
from fl.config import Config


class IPFSManager:
    def __init__(self, config: Config):
        self.config = config
        self.ipfs_client = ipfshttpclient.connect(f'/ip4/{self.config.IPFS_HOST}/tcp/{self.config.IPFS_PORT}')
    def upload_to_ipfs(self, model_params):
        buffer = io.BytesIO()   # 字节流对象，用于临时存储模型参数的序列化数据
        torch.save(model_params, buffer) # 将模型参数序列化并写入 buffer 中
        buffer.seek(0) # 将缓存区的当前位置设置为起始位置
        serialized_params = buffer.read() # 读取序列化后的字节数据
        res = self.ipfs_client.add_bytes(serialized_params) # 上传到ipfs
        return res

    def download_from_ipfs(self, ipfs_hash):
        serialized_params = self.ipfs_client.cat(ipfs_hash) # 下载字节数据
        buffer = io.BytesIO(serialized_params) # 写入缓存区
        return torch.load(buffer) # 加载并反序列化模型参顺，同时还原为 pytorch 张量对象
